"""
按比例拆分文件
"""
import os
import shutil
import random


def split_percent(root_dir, out_dir, file_type, train_ratio=0.8):
    png_files = os.listdir(root_dir)
    random.shuffle(png_files)

    # Compute the split indices
    total_files = len(png_files)
    train_end = int(total_files * train_ratio)

    # Split the files into train, val, and test
    train_files = png_files[:train_end]
    val_files = png_files[train_end:]

    def move_files(files, output_dir):
        for file in files:
            png_src_path = os.path.join(root_dir, file)
            shutil.copy(png_src_path, output_dir)

    # Move the files
    move_files(train_files, f"{out_dir}/train/{file_type}/")
    move_files(val_files, f"{out_dir}/val/{file_type}/")


if __name__ == '__main__':
    root_dir = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/fold1/cls_40_80_/"
    out_dir = "/media/hsmy/wanghao_18T/dataset/MIDOG2021/fold1/cls_40_80/"
    os.makedirs(out_dir + "train/p", exist_ok=True)
    os.makedirs(out_dir + "train/n", exist_ok=True)
    os.makedirs(out_dir + "val/p", exist_ok=True)
    os.makedirs(out_dir + "val/n", exist_ok=True)

    split_percent(root_dir + "p", out_dir, "p", train_ratio=0.8)
    split_percent(root_dir + "n", out_dir, "n", train_ratio=0.8)


